#include "Optimization.h"
#include "Exception.h"
#include "TransformInvariant.h"

namespace Compiler
{
	namespace Intermediate
	{
		class DeadCodeOptimizer : public IntraProcOptimizer
		{
		private:
			bool IsCritical(Instruction & instr)
			{
				return (instr.Func == Operation::Ret ||
						instr.Func == Operation::Call ||
						instr.Func == Operation::Store);
			}
			ControlFlowNode * FindNeareastMarkedPostDom(ControlFlowNode * node, IntSet & nodeMarks)
			{
				auto curNode = node->ReverseImmediateDominator;
				while (!nodeMarks.Contains(curNode->Id))
				{
					curNode = curNode->ReverseImmediateDominator;
				}
				return curNode;
			}
			void CompactNodes(ControlFlowGraph * graph)
			{
				List<ControlFlowNode*> traverse = graph->GetPostOrder();
				IntSet nodeMarks;
				nodeMarks.SetMax(graph->Nodes.Count());
				int count = 0;
				for (auto & node:traverse)
				{
					count ++;
					nodeMarks.Add(node->Id);
				}
				for (auto & node:traverse)
				{
					if (nodeMarks.Contains(node->Id))
					{
						List<ControlFlowNode*> newEntries;
						for (int i = 0; i<node->Entries.Count(); i++)
							if (nodeMarks.Contains(node->Entries[i]->Id))
								newEntries.Add(node->Entries[i]);
						
						node->Entries = _Move(newEntries);
						for (int i = 0; i<2; i++)
						{
							if (node->Exits[0])
							{
								if (!nodeMarks.Contains(node->Exits[0]->Id))
									node->Exits[0] = node->Exits[1];
							}
							else
								break;
						}
						if (node->GetExitCount()==1)
							node->Code.Last().Func = Operation::Jump;
					}
				}
				List<RefPtr<ControlFlowNode>> newNodes;
				newNodes.Reserve(count);
				for (int i = 0; i<graph->Nodes.Count(); i++)
				{
					if (nodeMarks.Contains(i))
					{
						graph->Nodes[i]->Id = newNodes.Count();
						newNodes.Add(graph->Nodes[i]);
					}
				}
				graph->Nodes = _Move(newNodes);

			}
			void Clean(ControlFlowGraph * graph)
			{
				CompactNodes(graph);

				bool changed = true;
				while (changed)
				{
					changed = false;
					List<ControlFlowNode*> traverse = graph->GetPostOrder();
					for (auto & node : traverse)
					{
						if (node == 0)
							continue;
						if (node == graph->Source || node == graph->Sink)
							continue;
						// Step 1: Fold redundant branch
						if (node->GetExitCount() == 2 && node->Exits[0] == node->Exits[1])
						{
							node->Exits[1] = 0;
							if (node->Code.Count())
							{
								node->Code.Last().Func = Operation::Jump;
							}
							node->Exits[0]->Entries.Remove(node);
							node->Exits[0]->Entries.Remove(node);
							node->Exits[0]->Entries.Add(node);
							changed = true;
						}
						if (node->GetExitCount() == 1)
						{
							// Step 2: Remove empty block
							if (node->Code.Count()==0 || node->Code.Count()==1 && node->Code.First().Func == Operation::Jump)
							{
								int entryId = node->Exits[0]->Entries.IndexOf(node);
								node->Exits[0]->Entries.RemoveAt(entryId);
								node->Exits[0]->Entries.InsertRange(entryId, node->Entries);
								for (auto & upperNode : node->Entries)
								{
									for (int j = 0; j<2; j++)
										if (upperNode->Exits[j] == node)
											upperNode->Exits[j] = node->Exits[0];
								}
								graph->Nodes[node->Id] = 0;
								node = 0;
								changed = true;
								continue;
							}
							// Step 3: merge block
							if (node->Exits[0] != graph->Sink && node->Exits[0]->Entries.Count() == 1)
							{
								auto exitNode = node->Exits[0];
								node->Code.LastNode()->Delete();
								for (auto & instr : exitNode->Code)
									node->Code.AddLast(instr);
								node->Exits[0] = exitNode->Exits[0];
								node->Exits[1] = exitNode->Exits[1];
								for (auto child : node->Exits)
									if (child)
									{
										int entryId = child->Entries.IndexOf(exitNode);
										child->Entries[entryId] = node;
									}
								graph->Nodes[exitNode->Id] = 0;
								int idx = traverse.IndexOf(exitNode);
								traverse[idx] = 0;
								changed = true;
								continue;
							}
							// Step 4: hoist branch
							if (node->Exits[0]->Code.Count() == 1 && node->Exits[0]->GetExitCount() == 2)
							{
								auto exitNode = node->Exits[0];
								node->Code.Last() = exitNode->Code.Last();
								for (int k = 0; k<2; k++)
								{
									node->Exits[k] = exitNode->Exits[k];
									int entryId = exitNode->Exits[k]->Entries.IndexOf(exitNode);
									exitNode->Exits[k]->Entries[entryId] = node;
								}
								changed = true;
							}
						}
					}
				}
				CompactNodes(graph);
				
				for (auto & node : graph->Nodes)
				{
					TransformInvariant::PlacePhiAtTop(node.Ptr());
					for (auto & instr : node->Code)
						instr.CFG_Node = node.Ptr();
				}
				graph->ComputeDominatorTree();
			}

			void RemoveEntry(ControlFlowNode * node, int id)
			{
				node->Entries.RemoveAt(id);
				for (auto & instr : node->Code)
					if (instr.Func == Operation::Phi)
						instr.Operands.RemoveAt(id);
			}
		public:
			DeadCodeOptimizer()
			{}
			virtual ProgramOptimizationResult Optimize(RefPtr<ControlFlowGraph> program) override
			{
				ProgramOptimizationResult result;
				result.Program = program;
				result.Changed = false;
				List<Instruction *> workList;
				IntSet nodeMark;
				nodeMark.SetMax(program->Nodes.Count());
				auto traverse = program->GetPostOrder();
				for (int i = traverse.Count()-1; i>=0; i--)
				{
					int instrId = 0;
					for (auto & instr : traverse[i]->Code)
					{
						if (IsCritical(instr))
						{
							nodeMark.Add(instr.CFG_Node->Id);
							instr.Mark = 1;
							workList.Add(&instr);
						}
						else
							instr.Mark = 0;
						instrId++;
					}
				}
				while (workList.Count())
				{
					auto &instr = *(workList.Last());
					workList.RemoveAt(workList.Count()-1);
					for (auto & op : instr.Operands)
					{
						if (op.IsVariable() && program->VarDefs[op.Var->Id] != 0)
						{
							auto defInstr = &program->VarDefs[op.Var->Id]->Value;
							if (defInstr->Mark == 0)
							{
								nodeMark.Add(defInstr->CFG_Node->Id);
								defInstr->Mark = 1;
								workList.Add(defInstr);
							}
						}
					}
					for (auto node : instr.CFG_Node->ReverseDominateFrontier)
					{
						/*if (node->Code.Count() == 0 ||
							node->Code.Last().Func != Operation::JumpF &&
							node->Code.Last().Func != Operation::JumpT)
							throw InvalidProgramException(L"Inconsistant CFG");*/
						auto &bInstr = node->Code.Last();
						if (bInstr.Mark == 0)
						{
							bInstr.Mark = 1;
							nodeMark.Add(bInstr.CFG_Node->Id);
							workList.Add(&bInstr);
						}
					}
				}
				for (int i = 0; i<program->Nodes.Count(); i++)
				{
					int instrId = 0;
					auto & code = program->Nodes[i]->Code;
					for (auto & iter = code.begin(); iter != code.end(); iter++)
					{
						auto & instr = iter.Current->Value;
						if (instr.Mark == 0)
						{
							if (instr.Func == Operation::Branch)
							{
								result.Changed = true;
								instr.Func = Operation::Jump;
								// find nearest marked postdominator
								auto postDom = FindNeareastMarkedPostDom(instr.CFG_Node, nodeMark);
								for (int j = 0; j<2; j++)
									if (instr.CFG_Node->Exits[j] != 0)
										instr.CFG_Node->Exits[j]->Entries.Remove(instr.CFG_Node);
								instr.CFG_Node->Exits[0] = postDom;
								instr.CFG_Node->Exits[1] = 0;
								postDom->Entries.Add(instr.CFG_Node);
							}
							else if (instr.Func != Operation::Jump)
							{
								result.Changed = true;
								iter.Current->Delete();
							}
						}
					}
				}
				/*if (result.Changed)
					Clean(result.Program.Ptr());*/
				
				return result;
			}
		};
		
		IntraProcOptimizer * CreateDeadCodeOptimizer()
		{
			return new DeadCodeOptimizer();
		}
	}
}